import vtk
import numpy as np

def meshplot(nodes, elements, scalars=None, range=None, show_lines=True, colormap=None, title=''):
    """
    使用VTK可视化有限元网格,并可选显示标量场。

    参数说明：
        nodes: (n,3) 的 numpy 数组，包含每个节点的 (x, y, z) 坐标。
        elements: 元素数据的字典列表，每个字典包含以下字段：
            - elements[i]['NodeIndices']: 第i个单元的节点索引(单元的连通性）。
            - elements[i]['CellType']: VTK单元类型(如 vtk.VTK_QUAD, vtk.VTK_TETRA 等）。
        scalars: 长度为 n 的 numpy 数组，包含节点的标量场值(可选）。
        range: 标量条的范围或极值区间。如果未指定，则使用全局最小最大值 [min(scalars), max(scalars)]。
        show_lines: 是否显示网格线(单元边界），默认为 True。
        colormap: 可选的标量场配色方案，默认使用 'viridis'。
        title: 可选的标量条标题。
    """
    # Create renderer
    renderer = vtk.vtkRenderer()
    renderer.GradientBackgroundOn()
    renderer.SetBackground(.7, .8, .9)
    renderer.SetBackground2(.1, .2, .3)

    # Create render window
    render_window = vtk.vtkRenderWindow()
    render_window.AddRenderer(renderer)

    # Create interactor style
    interactor_style = vtk.vtkInteractorStyleTrackballCamera()

    # Create interactor
    interactor = vtk.vtkRenderWindowInteractor()
    interactor.SetInteractorStyle(interactor_style)
    interactor.SetRenderWindow(render_window)

    # Create axes actor
    axes_actor = vtk.vtkAxesActor()
    axes_actor.SetShaftTypeToCylinder()
    axes_actor.SetTipTypeToCone()
    axes_actor.SetCylinderResolution(32)
    axes_actor.SetConeResolution(32)
    axes_actor.SetNormalizedLabelPosition(1.2, 1.2, 1.2)

    # Create orientation marker widget
    orientation_marker_widget = vtk.vtkOrientationMarkerWidget()
    orientation_marker_widget.SetViewport(0.0, 0.0, 0.22, 0.33)
    orientation_marker_widget.SetOrientationMarker(axes_actor)
    orientation_marker_widget.SetInteractor(interactor)
    orientation_marker_widget.EnabledOn()
    orientation_marker_widget.InteractiveOff()

    # Create unstructured grid
    unstructured_grid = vtk.vtkUnstructuredGrid()
    set_points(unstructured_grid, nodes)
    set_cells(unstructured_grid, elements)

    # Create data set mapper
    data_set_mapper = vtk.vtkDataSetMapper()
    data_set_mapper.SetInputData(unstructured_grid)
    data_set_mapper.Update()

    # Create unstructured grid actor
    unstructured_grid_actor = vtk.vtkActor()
    unstructured_grid_actor.SetMapper(data_set_mapper)
    unstructured_grid_actor.GetProperty().SetEdgeVisibility(show_lines)

    # Add unstructured grid actor
    renderer.AddActor(unstructured_grid_actor)

    # Include a nodal scalar field
    if scalars is not None:
        set_scalars(unstructured_grid, scalars)
        lookup_table = build_lookup_table(colormap)
        data_set_mapper.InterpolateScalarsBeforeMappingOn()
        data_set_mapper.SetLookupTable(lookup_table)
        data_set_mapper.ScalarVisibilityOn()
        if range is None:
            data_set_mapper.SetScalarRange(np.min(scalars), np.max(scalars))
        else:
            data_set_mapper.SetScalarRange(range[0], range[1])

        # Create scalar bar actor
        scalar_bar_actor = vtk.vtkScalarBarActor()
        scalar_bar_actor.SetLookupTable(lookup_table)
        scalar_bar_actor.SetTitle(title)
        scalar_bar_actor.SetLabelFormat("%+0.3e")
        scalar_bar_actor.SetNumberOfLabels(lookup_table.GetNumberOfColors() + 1)

        # Add scalar bar actor
        renderer.AddActor(scalar_bar_actor)

    # Start
    renderer.ResetCamera()
    interactor.Initialize()
    interactor.Start()


def set_points(unstructured_grid, point_data):
    """
    Utility function to set the points of an unstructured grid.
    """
    points = vtk.vtkPoints()
    for i, point in enumerate(point_data):
        points.InsertNextPoint(point[0], point[1], point[2])
    unstructured_grid.SetPoints(points)


def set_cells(unstructured_grid, cell_data):
    """
    Utility function to set the cells of an unstructured grid.
    """
    for cell in cell_data:
        cell_type = cell['CellType']
        node_indices = cell['NodeIndices']
        id_list = vtk.vtkIdList()
        for index in node_indices:
            id_list.InsertNextId(index - 1)
        unstructured_grid.InsertNextCell(cell_type, id_list)


def set_scalars(unstructured_grid, values):
    """
    Utility function to set the point scalars of an unstructured grid.
    """
    point_scalars = vtk.vtkDoubleArray()
    point_scalars.SetNumberOfValues(len(values))
    for i, value in enumerate(values):
        point_scalars.SetValue(i, value)
    unstructured_grid.GetPointData().SetScalars(point_scalars)


def build_lookup_table(cmap=None):
    """
    Utility function to build a lookup table.
    """
    if cmap is None:
        cmap = vtk.vtkColorTransferFunction()
        cmap.AddRGBPoint(0.0, 0.237, 0.299, 0.754)  # Default color map (viridis)
        cmap.AddRGBPoint(0.5, 0.831, 0.5, 0.0)
        cmap.AddRGBPoint(1.0, 0.8, 0.906, 0.30)

    lookup_table = vtk.vtkLookupTable()
    lookup_table.SetNumberOfColors(12)
    for i in range(12):
        color = cmap.GetColor(i / 11.0)
        lookup_table.SetTableValue(i, color[0], color[1], color[2], 1.0)
    lookup_table.Build()
    return lookup_table


# Example usage:
def Vtkexample():
    # Example nodes and elements
    nodes = np.array([
        [0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0],
        [1.0, 1.0, 0.0],
        [0.0, 1.0, 0.0]
    ])
    elements = [
        {'NodeIndices': [1, 2, 3, 4], 'CellType': vtk.VTK_QUAD}
    ]
    scalars = np.array([0.0, 0.5, 1.0, 0.75])
    meshplot(nodes, elements, scalars=scalars, range=[0, 1], show_lines=True, colormap=None, title="Scalar Field")